"""
RAG Embeddings - 大模型问答模块

该模块实现基于RAG的问答系统：
- 封装Ollama本地大模型API调用
- 根据用户问题检索相关文档片段
- 结合检索结果生成准确答案
- 支持批量问答和交互式问答

Author: xiaodeme@163.com
Date: 2025-07-09 18:50:30
"""

import requests
import json
from config import Config
from vector_db import VectorDB


class RAGQASystem:
    """基于RAG的问答系统"""
    
    def __init__(self, base_url=None, model=None):
        """
        初始化问答系统
        
        Args:
            base_url (str, optional): Ollama服务地址，默认使用配置文件设置
            model (str, optional): 使用的大模型名称，默认使用配置文件设置
        """
        self.base_url = base_url or Config.OLLAMA_BASE_URL
        self.model = model or Config.OLLAMA_MODEL
        
        # 初始化向量数据库
        print("[INFO] 初始化向量数据库连接...")
        self.vector_db = VectorDB()
        print("[INFO] RAG问答系统初始化完成")

    def _format_context(self, query_results, max_context_length=2000):
        """
        格式化检索到的上下文
        
        Args:
            query_results (dict): 向量数据库查询结果
            max_context_length (int): 最大上下文长度，避免超出模型限制
            
        Returns:
            str: 格式化后的上下文字符串
        """
        documents = query_results.get('documents', [[]])
        distances = query_results.get('distances', [[]])
        
        if not documents or not documents[0]:
            return "未找到相关文档内容"
        
        # 组合文档和相似度分数
        context_parts = []
        for i, (doc, distance) in enumerate(zip(documents[0], distances[0] if distances else [])):
            similarity = 1 - distance if distance is not None else 0.0
            context_parts.append(f"文档片段{i+1}（相似度：{similarity:.3f}）:\n{doc}\n")
        
        # 拼接上下文，如果超长则截断
        full_context = "\n".join(context_parts)
        if len(full_context) > max_context_length:
            full_context = full_context[:max_context_length] + "...[内容被截断]"
            
        return full_context

    def get_llm_response(self, query, context, timeout=30):
        """
        调用Ollama API获取大模型回答
        
        Args:
            query (str): 用户问题
            context (str): 检索到的相关上下文
            timeout (int): 请求超时时间（秒）
            
        Returns:
            str: 大模型生成的回答
        """
        # 构建优化的提示词
        prompt = f"""请基于以下提供的上下文信息回答用户问题。如果上下文中没有相关信息，请明确说明。

上下文信息：
{context}

用户问题：{query}

请提供准确、详细的回答："""

        # 构建API请求数据
        url = f"{self.base_url}/api/chat"
        headers = {"Content-Type": "application/json"}
        data = {
            "model": self.model,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "stream": False,  # 不使用流式响应
            "options": {
                "temperature": 0.1,  # 降低随机性，提高回答准确性
                "top_p": 0.9
            }
        }

        try:
            print(f"[INFO] 正在调用大模型 {self.model} 生成回答...")
            response = requests.post(url, headers=headers, json=data, timeout=timeout)
            response.raise_for_status()

            # 解析响应
            result = response.json()
            answer = result.get("message", {}).get("content", "").strip()
            
            if not answer:
                return "抱歉，模型没有返回有效回答。"
                
            print(f"[SUCCESS] 大模型回答生成完成")
            return answer
            
        except requests.exceptions.Timeout:
            return f"请求超时（{timeout}秒），请检查模型服务是否正常或增加超时时间。"
        except requests.exceptions.ConnectionError:
            return f"无法连接到Ollama服务（{self.base_url}），请检查服务是否启动。"
        except requests.exceptions.HTTPError as e:
            return f"HTTP错误：{e}，请检查模型名称是否正确。"
        except json.JSONDecodeError:
            return "服务器返回的响应格式错误。"
        except Exception as e:
            return f"调用大模型API时发生错误：{e}"

    def answer_question(self, question, n_results=3):
        """
        回答单个问题
        
        Args:
            question (str): 用户问题
            n_results (int): 检索文档数量
            
        Returns:
            dict: 包含问题、上下文、答案的字典
        """
        print(f"\n{'='*60}")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        # Step 1: 检索相关文档
        print("[1/3] 检索相关文档...")
        query_results = self.vector_db.query(question, n_results=n_results)
        
        # Step 2: 格式化上下文
        print("[2/3] 处理上下文信息...")
        context = self._format_context(query_results)
        
        # Step 3: 生成答案
        print("[3/3] 生成模型回答...")
        answer = self.get_llm_response(question, context)
        
        print(f"\n回答:")
        print(f"{answer}")
        print(f"\n{'='*60}")
        
        return {
            "question": question,
            "context": context,
            "answer": answer
        }

    def batch_qa(self, questions, n_results=3):
        """
        批量问答
        
        Args:
            questions (list): 问题列表
            n_results (int): 每个问题检索的文档数量
            
        Returns:
            list: 问答结果列表
        """
        print(f"\n开始批量问答，共 {len(questions)} 个问题...")
        results = []
        
        for i, question in enumerate(questions, 1):
            print(f"\n处理问题 {i}/{len(questions)}")
            result = self.answer_question(question, n_results)
            results.append(result)
            
        print(f"\n批量问答完成，共处理 {len(results)} 个问题")
        return results

    def interactive_qa(self):
        """交互式问答模式"""
        print("\n进入交互式问答模式")
        print("输入 'quit' 或 'exit' 退出")
        print("=" * 50)
        
        while True:
            try:
                question = input("\n💬 请输入您的问题: ").strip()
                
                if question.lower() in ['quit', 'exit', '退出']:
                    print("👋 感谢使用RAG问答系统！")
                    break
                    
                if not question:
                    print("⚠️  请输入有效问题")
                    continue
                    
                self.answer_question(question)
                
            except KeyboardInterrupt:
                print("\n\n👋 用户中断，退出程序")
                break
            except Exception as e:
                print(f"处理问题时发生错误: {e}")


def main():
    """主函数：演示RAG问答系统功能"""
    
    try:
        # 初始化问答系统
        qa_system = RAGQASystem()
        
        # 检查向量数据库是否有数据
        info = qa_system.vector_db.get_collection_info()
        if info['document_count'] == 0:
            print(" 向量数据库为空，请先运行 vector_db.py 添加文档")
            return
        
        # 定义测试问题
        test_questions = [
            "市场准入可以包括哪些方面？",
            "市场准入需要搭建哪些团队？",
            "体系有效运行和上市后监管的重点是什么？",
            "软件类产品如何选择检测机构？",
            "临床试验大致可以分为几个阶段？"
        ]
        
        print("RAG问答系统演示")
        print("1. 批量问答演示")
        print("2. 交互式问答")
        
        choice = input("\n请选择模式 (1/2): ").strip()
        
        if choice == "1":
            # 批量问答演示
            qa_system.batch_qa(test_questions)
        elif choice == "2":
            # 交互式问答
            qa_system.interactive_qa()
        else:
            print("无效选择，执行默认的批量问答演示")
            qa_system.batch_qa(test_questions)
            
    except Exception as e:
        print(f"[ERROR] 系统初始化失败: {e}")
        print("请确保Ollama服务正在运行并已加载指定模型")


if __name__ == '__main__':
    main()